import json
import numpy as np
import pandas as pd
from keras_bert import Tokenizer, load_vocabulary, load_trained_model_from_checkpoint, get_checkpoint_paths
import tensorflow as tf
import keras
import keras.backend as K
from keras.preprocessing import sequence
from keras.layers import Dense, Input, Masking, Concatenate, Lambda, Multiply, Activation
from keras import Model
import jieba
import editdistance
import os
import matplotlib.pyplot as plt
from read_data import get_data_id
get_data_filename = '../../data/train_deal_data.json'
evaluate_data, evaluate_question = get_data_id(get_data_filename, methods='evaluate')
cel_num_list = []
for i in range(len(evaluate_data[3])):
    cel_num = 0
    for j in range(len(evaluate_data[3][0])):
        if evaluate_data[3][0][j] == 1:
            cel_num += 1
    cel_num_list.append(cel_num)

def get_sql(predict_data, question, cel_nums):  # questions,cel_nums可以从输入的数据中获得，perdict_data可以从输入的数据中predict获得
    p_cond_conn_op = predict_data[0]  # (4, 3)
    p_sel_agg = predict_data[1]       # (4, sel_len, 7)
    p_cond_cel = predict_data[2]      # (4, 412, 412)
    value_op = predict_data[3]        # (4, 412 , 5)

    #
    p_cond_conn_op_list = np.argmax(p_cond_conn_op, axis=-1)[0]

    p_sel_agg_list = []
    sel_agg_max = np.argmax(p_sel_agg, axis=-1)  # (4, 412, 1)
    print("sel_agg_max", sel_agg_max)
    for i in range(len(sel_agg_max)):
        for cel, agg in enumerate(sel_agg_max[i][:cel_nums[i]]): # (412, 1)
            print(cel, agg)
            if agg != 6:  # inputs_data[5] headers
                p_sel_agg_list.append([cel, agg])
        if len(p_sel_agg_list) == 0:
            return None
    print(p_sel_agg_list)


    # 设置一个预测conds中的条件个数
    sql = []
    v_start_all_list, v_str_all_list, ops_all_list = [], [], []
    # (4, 412 , 5)
    for j in range(len(value_op)):
        v_str_len = 0
        v_start = 0
        # get the vlaue's start and len
        v_str_len_list, v_start_list, v_str_list = [], [], []
        ops = {}
        for index, values in enumerate(value_op[j].argmax(axis=-1)):
            if values != 4:
                if v_start == 0:
                    v_start = index
                    v_start_list.append(v_start)
                ops[index] = values
                v_str_len += 1
            else:
                v_start = 0
                v_str_len_list.append(v_str_len)
                v_str_len = 0
        for i in v_str_len_list:
            if i != 0:
                v_str_list.append(i)
        v_start_all_list.append(v_start_list)
        v_str_all_list.append(v_str_list)
        ops_all_list.append(ops)
        # [9, 13] [3, 4] {9: 2, 10: 2, 11: 2, 13: 2, 14: 2, 15: 2, 16: 2}
        print(v_start_all_list, v_str_all_list, ops_all_list)

    conds = []
    for j in range(len(v_start_all_list)):  # 4个
        for i in range(len(v_start_all_list[j])):  # 操作一个
            v_start = v_start_all_list[j][i]
            v_end = v_start_all_list[j][i] + v_str_all_list[j][i]
            str = question[j][v_start - 1: v_end - 1]
            op = ops_all_list[j][v_start]
            print(f"p_cond_cel.shape is {p_cond_cel.shape}")
            print(f"p_cond_cel[0][v_start: v_end][:].shape is {p_cond_cel[0][v_start: v_end][:].shape}")
            cel = np.mean(p_cond_cel[0][v_start: v_end][:], axis=0).argmax() - 1  # p_cond_cel[0][v_start: v_end][:].shape (4, 412)
            print(f"cel is {np.mean(p_cond_cel[0][v_start: v_end][:], axis=0)}")
            conds.append([cel, op, str])
        sql.append({"agg":p_sel_agg_list[j][1],"cond_conn_op": p_cond_conn_op_list[j], "sel": p_sel_agg_list[j][0], "conds":conds})
    return sql

def is_equal(sql, pred_sql):
    return sql['agg']==pred_sql['agg'] and sql['cond_conn_op']==pred_sql['cond_conn_op'] and \
            sql['sel'] == pred_sql['sel'] and sql['conds']==pred_sql['conds']

'''
[9, 13] [3, 4] {9: 2, 10: 2, 11: 2, 13: 2, 14: 2, 15: 2, 16: 2}
p_cond_cel.shape is (1, 412, 412)
p_cond_cel[0][v_start: v_end][:].shape is (3, 412)
cel is 0
p_cond_cel.shape is (1, 412, 412)
p_cond_cel[0][v_start: v_end][:].shape is (4, 412)
cel is 0
{'agg': 5, 'cond_conn_op': 2, 'sel': 2, 'conds': [[-1, 2, '大黄蜂'], [-1, 2, '密室逃生']]}
'''
# ('or', [(0, 2, '大黄蜂'), (0, 2, '密室逃生')])
# "sql": {"agg": [5], "cond_conn_op": 2, "sel": [2], "conds": [[0, 2, "大黄蜂"], [0, 2, "密室逃生"]]}
# sql = {"agg": [5], "cond_conn_op": 2, "sel": [2], "conds": [[0, 2, "大黄蜂"], [0, 2, "密室逃生"]]}
batch_size = 4
print(get_sql(evaluate_data[0:batch_size], evaluate_question[0:batch_size], cel_num_list[0:batch_size]))

